from model import *
from utils import *
import random

mnist_test = torchvision.datasets.MNIST(root="./data",
                                        train=False,
                                        transform=transforms.ToTensor(),
                                        download=True)

model = CVAE()
model.load_state_dict(torch.load('cvae.params'))
model.eval()

figsize = (20, 10)
_, axes = plt.subplots(5, 10, figsize=figsize)
axes = axes.flatten()

for i in range(50):
    t = i % 10
    x, l = torch.randn(10).unsqueeze(0).to(device), torch.zeros(
        1, 10).scatter(1, torch.tensor([t]).unsqueeze(1), 1).to(device)
    img = model.decoder(x, l).reshape(28, 28).cpu().detach().numpy()
    axes[i].imshow(img)
    axes[i].axes.get_xaxis().set_visible(False)
    axes[i].axes.get_yaxis().set_visible(False)

plt.show()
